#include "Sessions.h"

bool Sessions::DefineMap() {
	SessionTypesByHost.insert(make_pair("www.guap.ru", BROWSING));
	SessionTypesByHost.insert(make_pair("guap.ru", BROWSING));
	SessionTypesByHost.insert(make_pair("fs.guap.ru", VIDEO));
	SessionTypesByHost.insert(make_pair("ivi.ru", VIDEO));
	SessionTypesByHost.insert(make_pair("smotretvideo.com", VIDEO));
	SessionTypesByHost.insert(make_pair("smotri.com", VIDEO));
	SessionTypesByHost.insert(make_pair("kinobar.net", VIDEO));

	SessionTypesByContent.insert(make_pair("video", VIDEO));
	SessionTypesByContent.insert(make_pair("image/gif", BROWSING));
	SessionTypesByContent.insert(make_pair("image/jpeg", BROWSING));
	SessionTypesByContent.insert(make_pair("image/x-icon", BROWSING));
	SessionTypesByContent.insert(make_pair("text/html", BROWSING));

	return true;
}

bool Sessions::IpCompare(u_char neededIP[], const u_char currentIP[])
{
    bool flag = true;
    for(int i = 0; i < 4; i++) {
        if(neededIP[i] != currentIP[i]) {
            flag = false;
            break;
        }
    }
    return flag;
}

bool Sessions::TcpKeyCompare(TcpKey firstKey, TcpKey secondKey) {
	if(IpCompare(firstKey.srcIP,secondKey.srcIP) && IpCompare(firstKey.dstIP,secondKey.dstIP) && (firstKey.srcPort == secondKey.srcPort) && (firstKey.dstPort == secondKey.dstPort) && (firstKey.protocol == secondKey.protocol)) {
		return true;
	}
	return false;
}

int  Sessions::WriteList(string resFile, list<TcpPacket> listToWrite) {

    FILE *file = fopen(resFile.c_str(), "a+");

    int bytes_written = 0;

    for (list<TcpPacket>::iterator it = listToWrite.begin(); it != listToWrite.end(); ++it)
    {
            bytes_written = fwrite(it->payload, sizeof(u_char), it->packetLen - 66, file);
            bytes_written = fwrite("\n\n\n\n", sizeof(char), sizeof("\n\n\n\n"), file);
    }

    fclose(file);

    return bytes_written;
}

bool Sessions::PrintType(int type){
	if(type == 1) cout << "VIDEO\n";
	if(type == 2) cout << "BROWSING\n";
	if(type == 3) cout << "DOWNLOAD\n";
	if(type == 0) cout << "NON_CLASSIFIED\n";

	return true;
}

bool Sessions::PrintMonitor(map<TcpKey, list<TcpPacket> >::iterator itPrint, TcpPacket tempPacket){
	if (FindHttp(itPrint->second) == 1) {
		for(int i = 0; i < 4; i++) {
			printf("%d", itPrint->first.srcIP[i]);
			cout << '.';
		}
		cout << " - ";
		for(int i = 0; i < 4; i++) {
			printf("%d", itPrint->first.dstIP[i]);
			cout << '.';
		}

		cout << "   " << "HTTP  ";

		int type = 0;
		pair<TcpKey,list<TcpPacket> > elem = *itPrint;

		if(itPrint->first.link == "uplink") {
			type = FindHttpHostName(tempPacket);
			if (type != 0) {

				PrintType(type);
				cout << "\n\n";

				for (map<TcpKey, list<TcpPacket> >::iterator itRev = AllSessions.begin(); itRev != AllSessions.end(); ++itRev) //iterate through list of sessions to find session or to add new
				{

					if(IpCompare(elem.first.dstIP, itRev->first.srcIP) && IpCompare(elem.first.srcIP, itRev->first.dstIP)) {
						UpdateMapKey(itRev, type);
						break;
					}
				}
				UpdateMapKey(itPrint, type);
			}
			else {
				PrintType(itPrint->first.sessType);
				cout << "\n\n";
			}
		}
		else if(itPrint->first.link == "downlink") {
			type = FindHttpContentType(tempPacket);
			if (type != 0) {

				PrintType(type);
				cout << "\n\n";

				for (map<TcpKey, list<TcpPacket> >::iterator itRev = AllSessions.begin(); itRev != AllSessions.end(); ++itRev) //iterate through list of sessions to find session or to add new
				{

					if(IpCompare(elem.first.dstIP, itRev->first.srcIP) && IpCompare(elem.first.srcIP, itRev->first.dstIP)) {
						UpdateMapKey(itRev, type);
						break;
					}
				}
				UpdateMapKey(itPrint, type);
			}
			else {
				PrintType(itPrint->first.sessType);
				cout << "\n\n";
			}
		}
	}

	else {
		for(int i = 0; i < 4; i++) {
			printf("%d", itPrint->first.srcIP[i]);
			cout << '.';
		}
		cout << " - ";
		for(int i = 0; i < 4; i++) {
			printf("%d", itPrint->first.dstIP[i]);
			cout << '.';
		}

		cout << "   " << "non HTTP" << "\n\n\n";
	}
	return true;
}

bool Sessions::UpdateMapKey(map<TcpKey, list<TcpPacket> >::iterator it, int sessType) {

	struct TcpKey tempTcpKey;
	
	//Fullfill map key structure with the recieved packet info
	tempTcpKey.srcIP[0] = it->first.srcIP[0];
	tempTcpKey.srcIP[1] = it->first.srcIP[1];
	tempTcpKey.srcIP[2] = it->first.srcIP[2];
	tempTcpKey.srcIP[3] = it->first.srcIP[3];

	tempTcpKey.dstIP[0] = it->first.dstIP[0];
	tempTcpKey.dstIP[1] = it->first.dstIP[1];
	tempTcpKey.dstIP[2] = it->first.dstIP[2];
	tempTcpKey.dstIP[3] = it->first.dstIP[3];

	tempTcpKey.srcPort = it->first.srcPort;
	tempTcpKey.dstPort = it->first.dstPort;

	tempTcpKey.protocol = it->first.protocol;

	tempTcpKey.link = it->first.link;
	tempTcpKey.sessType = sessType;

	list<TcpPacket> tempList (it->second.begin(), it->second.end());

	AllSessions.erase(it);
	AllSessions.insert(make_pair(tempTcpKey,tempList));

	return 1;
}


int Sessions::FindHttpHostName(TcpPacket tempPacket){
	printf("Host  ");
	
	list<TcpPacket> tempList;
	tempList.push_front(tempPacket);
	u_char* hostName;

	if(FindHttp(tempList)) {
		int status = 0;
		int startPtr = 0;
		int endPtr = 0;

		for (int i = 0; i < tempPacket.packetLen-66; i++)
		{
			if(status == 0) {
				if(tempPacket.payload[i] == ('H')) {
					status = 1;
					continue;
				}
				else status = 0;
			}

			if(status == 1) {
				if(tempPacket.payload[i] == ('o')) {
					status = 2;
					continue;
				}
				else status = 0;
			}

			if(status == 2) {
				if(tempPacket.payload[i] == ('s')) {
					status = 3;
					continue;
				}
				else status = 0;
			}

			if(status == 3) {
				if(tempPacket.payload[i] == ('t')) {
					status = 4;
					continue;
				}
				else status = 0;
			}

			if(status == 4) {
				if(tempPacket.payload[i] == (':')) {
					startPtr = i + 2;
					while(tempPacket.payload[i] != '\r') {
						i++;
					}
					endPtr = i ;

					hostName = (u_char*)calloc(endPtr - startPtr + 1, sizeof(u_char));

					for (int i = 0; i < endPtr-startPtr; i++)
					{
						hostName[i] = tempPacket.payload[i+startPtr];
					}
					bool flag = true;

					string hostString((char*) hostName);
					cout << hostString << "  ";
					for (map<string, int>::iterator it = SessionTypesByHost.begin(); it != SessionTypesByHost.end(); ++it) //iterate through list of sessions to find session or to add new
					{
						flag = true;
						pair<string, int> elem = *it;
					
						if((int)elem.first.size() == endPtr - startPtr) {
							for(int i = 0; i < endPtr - startPtr; i ++) {
								if(elem.first[i] != hostString[i]){
									flag = false;
								}
							}
							if (flag == true)
								return elem.second;
						}
					}
				}
				else status = 0;
			}

		}
	}
	return 0;
}

int Sessions::FindHttpContentType(TcpPacket tempPacket){
	printf("Content  ");

	list<TcpPacket> tempList;
	tempList.push_front(tempPacket);
	u_char* contentType;

	if(FindHttp(tempList)) {
		int status = 0;
		int startPtr = 0;
		int endPtr = 0;

		for (int i = 0; i < tempPacket.packetLen-66; i++)
		{
			if(status == 0) {
				if(tempPacket.payload[i] == ('C')) {
					status = 1;
					continue;
				}
				else status = 0;
			}

			if(status == 1) {
				if(tempPacket.payload[i] == ('o')) {
					status = 2;
					continue;
				}
				else status = 0;
			}

			if(status == 2) {
				if(tempPacket.payload[i] == ('n')) {
					status = 3;
					continue;
				}
				else status = 0;
			}

			if(status == 3) {
				if(tempPacket.payload[i] == ('t')) {
					status = 4;
					continue;
				}
				else status = 0;
			}

			if(status == 4) {
				if(tempPacket.payload[i] == ('e')) {
					status = 5;
					continue;
				}
				else status = 0;
			}

			if(status == 5) {
				if(tempPacket.payload[i] == ('n')) {
					status = 6;
					continue;
				}
				else status = 0;
			}

			if(status == 6) {
				if(tempPacket.payload[i] == ('t')) {
					status = 7;
					continue;
				}
				else status = 0;
			}

			if(status == 7) {
				if(tempPacket.payload[i] == ('-')) {
					status = 8;
					continue;
				}
				else status = 0;
			}

			if(status == 8) {
				if(tempPacket.payload[i] == ('T')) {
					status = 9;
					continue;
				}
				else status = 0;
			}

			if(status == 9) {
				if(tempPacket.payload[i] == ('y')) {
					status = 10;
					continue;
				}
				else status = 0;
			}

			if(status == 10) {
				if(tempPacket.payload[i] == ('p')) {
					status = 11;
					continue;
				}
				else status = 0;
			}

			if(status == 11) {
				if(tempPacket.payload[i] == ('e')) {
					status = 12;
					continue;
				}
				else status = 0;
			}

			if(status == 12) {
				if(tempPacket.payload[i] == (':')) {
					startPtr = i + 2;
					while(tempPacket.payload[i] != '\r') {
						i++;
						if (tempPacket.payload[i] == ';') 
							break;
					}
					endPtr = i ;

					contentType = (u_char*)calloc(endPtr - startPtr + 1, sizeof(u_char));

					for (int i = 0; i < endPtr-startPtr; i++)
					{
						contentType[i] = tempPacket.payload[i+startPtr];
					}
					bool flag = true;

					string contentString((char*) contentType);
					cout << contentString << "  ";

		
					for (map<string, int>::iterator it = SessionTypesByContent.begin(); it != SessionTypesByContent.end(); ++it) //iterate through list of sessions to find session or to add new
					{
						flag = true;
						pair<string, int> elem = *it;

						if((int)elem.first.size() == endPtr - startPtr) {
							for(int i = 0; i < endPtr - startPtr; i ++) {
								if(elem.first[i] != contentString[i]){
									flag = false;
								}
							}
							if (flag == true)
								return elem.second;
						}
					}
				}
				else status = 0;
			}

		}
	}
	return 0;
}

bool Sessions::FindHttp(list<TcpPacket> listToSearch) {
	
	int status = 0;
	/*for (list<TcpPacket>::iterator it = listToSearch.begin(); it != listToSearch.end(); ++it){
		TcpPacket elem = *it;
		for (int i = 0; i < elem.packetLen-66; i++)
		{
			printf("%c", elem.payload[i]);
		}
		cout << "\n\n\n";
	}*/

	for (list<TcpPacket>::iterator it = listToSearch.begin(); it != listToSearch.end(); ++it){
		TcpPacket elem = *it;
		
		for (int i = 0; i < it->packetLen-66; i++)
		{
			if(status == 0) {
				if(elem.payload[i] == ('H')) {
					status = 1;
					continue;
				}
				else status = 0;
			}

			if(status == 1) {
				if(elem.payload[i] == ('T')) {
					status = 2;
					continue;
				}
				else status = 0;
			}

			if(status == 2) {
				if(elem.payload[i] == ('T')) {
					status = 3;
					continue;
				}
				else status = 0;
			}

			if(status == 3) {
				if(elem.payload[i] == ('P')) {
					return true;
				}
				else status = 0;
			}

		}
	}

	return false;
}

bool Sessions::AddPacket(const u_char *packet, u_char compIp[], bool searchHttp) {

    //check SYN ACK ack-10, syn-02, syn ack 12 fin ack 11 rst 04 and TCP protocol
    if((packet[47] == 0x10) || (packet[47] == 0x02) || (packet[47] == 0x12) || (packet[47] == 0x11) || (packet[47] == 0x04) || (packet[23] != 0x06)) {
      return false;
    }

    struct TcpKey tempTcpKey;
    struct TcpPacket tempPacket;

    //Fullfill packet structure with the recived packet info
    tempPacket.packetLen = (packet[16] << 8 | packet[17]) + 14;

    tempPacket.seqNumber = (unsigned(packet[38]) << 24) | (unsigned(packet[39]) << 16) | (unsigned(packet[40]) << 8) | (unsigned(packet[41]));

	tempPacket.payload = (u_char*)calloc(tempPacket.packetLen - 66, sizeof(u_char));

	for (int i = 0; i < tempPacket.packetLen-66; i++)
	{
		tempPacket.payload[i] = packet[i+66];
	}

    //Fullfill map key structure with the recieved packet info
    tempTcpKey.srcIP[0] = packet[26];
    tempTcpKey.srcIP[1] = packet[27];
    tempTcpKey.srcIP[2] = packet[28];
    tempTcpKey.srcIP[3] = packet[29];

    tempTcpKey.dstIP[0] = packet[30];
    tempTcpKey.dstIP[1] = packet[31];
    tempTcpKey.dstIP[2] = packet[32];
    tempTcpKey.dstIP[3] = packet[33];

    tempTcpKey.srcPort = packet[34] << 8 | packet[35];
    tempTcpKey.dstPort = packet[36] << 8 | packet[37];

    tempTcpKey.protocol = packet[23];

	if(IpCompare(compIp,tempTcpKey.srcIP) == true) tempTcpKey.link = "uplink";
	if(IpCompare(compIp,tempTcpKey.dstIP) == true && tempTcpKey.srcPort == 80) tempTcpKey.link = "downlink";

	tempTcpKey.sessType = 0;

	DefineMap();
    //look through the list of lists of packets and see that there is no the same session list
    if (AllSessions.size() == 0) { //no sessions at all create the first
        list<TcpPacket> tcpSession;

        tcpSession.push_front(tempPacket);

        AllSessions.insert(make_pair(tempTcpKey, tcpSession));

		if(searchHttp == 1) 
		{
			PrintMonitor(AllSessions.begin(), tempPacket);
		}
	}
    else {

        bool sessionFound = false;

        for (map<TcpKey, list<TcpPacket> >::iterator it = AllSessions.begin(); it != AllSessions.end(); ++it) //iterate through list of sessions to find session or to add new
        {
			pair<TcpKey,list<TcpPacket> > elem = *it;

            if(IpCompare(elem.first.srcIP, tempTcpKey.srcIP) == 0) continue;
			if(IpCompare(elem.first.dstIP, tempTcpKey.dstIP) == 0) continue;

            if(it->first.srcPort != tempTcpKey.srcPort) continue;
            if(it->first.dstPort != tempTcpKey.dstPort) continue;
            if(it->first.protocol != tempTcpKey.protocol) continue;

            sessionFound = true;

            bool packetAdded = false;

            for (list<TcpPacket>::reverse_iterator it2 = it->second.rbegin(); it2 != it->second.rend(); it2++) //itrate through list of packets in session to insert into the right place
            {
                if(tempPacket.seqNumber >= it2->seqNumber) {
                    it->second.insert(it2.base(), tempPacket);
					if(searchHttp == 1) {
						PrintMonitor(it, tempPacket);
					}

                    packetAdded = true;
                    break;
                }
            }

            if (packetAdded == true)
            {
                break;
            }

            it->second.push_front(tempPacket);
			if(searchHttp == 1) {
				PrintMonitor(it, tempPacket);
			}

            break;
        }

        if(sessionFound == false) {
            list<TcpPacket> tcpSession;

            tcpSession.push_front(tempPacket);

            AllSessions.insert(make_pair(tempTcpKey, tcpSession));
			
			if(searchHttp == 1) 
			{
				for (map<TcpKey, list<TcpPacket> >::iterator it = AllSessions.begin(); it != AllSessions.end(); ++it) //iterate through list of sessions to find session or to add new
				{
					pair<TcpKey,list<TcpPacket> > elem = *it;
					if(TcpKeyCompare(elem.first, tempTcpKey) == 1) {
						PrintMonitor(it, tempPacket);
						break;
					}
				}
					
			}
			
        }
    }

    return true;
}

bool Sessions::DumpSession(u_char* defIp, int packetNum, char* resFileName){
	for (map<TcpKey, list<TcpPacket> >::iterator it = AllSessions.begin(); it != AllSessions.end(); ++it) //iterate through list of sessions to find session or to add new
    {
		if(it->second.size() == packetNum && IpCompare(defIp, it->first.srcIP) && (it->first.link == "downlink")) {
			WriteList(resFileName, it->second);
			return true;
		}
	}
	return false;
}


bool Sessions::DumpReverseSession(u_char* defIp, char* resFileName){
	for (map<TcpKey, list<TcpPacket> >::iterator it = AllSessions.begin(); it != AllSessions.end(); ++it) //iterate through list of sessions to find session or to add new
    {
		if(IpCompare(defIp, it->first.dstIP) && (it->first.link == "uplink")) {
			WriteList(resFileName, it->second);
			return true;
		}
	}
	return false;
}
